import torch
import copy
import torch.nn as nn
from torch.nn import functional as F
from torch.nn import Module
from torch.nn import MultiheadAttention
from torch.nn import ModuleList
from torch.nn.init import xavier_uniform_
from torch.nn import Dropout
from torch.nn import Linear
from torch.nn import LayerNorm
import math
from torch.autograd import Variable


class Dual_lstm_cell(nn.Module):
    """
    双LSTM单元类，实现视觉和音频模态的LSTM单元，支持跨模态信息交互
    每个模态的LSTM门控计算可能受到另一模态的影响（当前代码中注释了跨模态项，可启用）
    """
    def __init__(self, visual_input_dim, audio_input_dim, hidden_dim, alph=0.5, bias=True):
        super(Dual_lstm_cell, self).__init__()

        # 保存输入维度、隐藏层维度和跨模态权重参数
        self.visual_input_dim = visual_input_dim  # 视觉输入特征维度
        self.audio_input_dim  = audio_input_dim   # 音频输入特征维度
        self.hidden_dim = hidden_dim              # LSTM隐藏状态维度
        self.alph = alph                          # 跨模态信息融合权重（0.5表示等权重融合）

        # 视觉模态LSTM门控的线性层（输入特征->门控，隐藏状态->门控）
        # 4*hidden_dim对应输入门、遗忘门、细胞门、输出门
        self.vs_linear = nn.Linear(self.visual_input_dim, 4 * self.hidden_dim, bias=bias)
        self.vh_linear = nn.Linear(self.hidden_dim, 4* self.hidden_dim, bias=bias)

        # 音频模态LSTM门控的线性层（第一套，用于基础计算）
        self.as_linear = nn.Linear(self.audio_input_dim, 4 * self.hidden_dim, bias=bias)
        self.ah_linear = nn.Linear(self.hidden_dim, 4 * self.hidden_dim, bias=bias)

        # 音频模态LSTM门控的第二套线性层（可能用于跨模态交互时的视觉信息输入）
        self.as_linear2 = nn.Linear(self.audio_input_dim, 4*self.hidden_dim, bias=bias)
        self.ah_linear2 = nn.Linear(self.hidden_dim, 4*self.hidden_dim, bias=bias)

        # 视觉模态LSTM门控的第二套线性层（可能用于跨模态交互时的音频信息输入）
        self.vs_linear2 = nn.Linear(self.visual_input_dim, 4*self.hidden_dim, bias=bias)
        self.vh_linear2 = nn.Linear(self.hidden_dim, 4*self.hidden_dim, bias=bias)

        # 初始化参数
        self.reset_parameters()

    def reset_parameters(self):
        """参数初始化：使用均匀分布初始化权重，范围为[-1/sqrt(hidden_dim), 1/sqrt(hidden_dim)]"""
        std = 1.0 / math.sqrt(self.hidden_dim)
        for w in self.parameters():
            w.data.uniform_(-std, std)

    def forward(self, visual_state, visual_hidden, visual_cell, audio_state, audio_hidden, audio_cell):
        """
        前向传播：计算单步LSTM的状态更新
        Args:
            visual_state: 视觉输入特征 (batch_size, visual_input_dim)
            visual_hidden: 视觉上一时刻隐藏状态 (batch_size, hidden_dim)
            visual_cell: 视觉上一时刻细胞状态 (batch_size, hidden_dim)
            audio_state: 音频输入特征 (batch_size, audio_input_dim)
            audio_hidden: 音频上一时刻隐藏状态 (batch_size, hidden_dim)
            audio_cell: 音频上一时刻细胞状态 (batch_size, hidden_dim)
        Returns:
            更新后的视觉隐藏状态、视觉细胞状态、音频隐藏状态、音频细胞状态
        """
        # 计算视觉模态的门控（当前注释了音频跨模态项，解开可加入音频对视觉的影响）
        visual_gates = self.vs_linear(visual_state) + self.vh_linear(visual_hidden)
            # 可选跨模态融合：+ self.alph*self.as_linear(audio_state) + self.alph*self.ah_linear(audio_hidden)

        # 计算音频模态的门控（当前注释了视觉跨模态项，解开可加入视觉对音频的影响）
        audio_gates = self.as_linear2(audio_state) + self.ah_linear2(audio_hidden)
            # 可选跨模态融合：+ self.alph*self.vs_linear2(visual_state) + self.alph*self.vh_linear2(visual_hidden)

        # 将门控分割为输入门(i)、遗忘门(f)、细胞门(c)、输出门(o)
        visual_i_gate, visual_f_gate, visual_c_gate, visual_o_gate = visual_gates.chunk(4, 1)
        audio_i_gate, audio_f_gate, audio_c_gate, audio_o_gate = audio_gates.chunk(4, 1)

        # 视觉模态门控激活函数：sigmoid用于输入/遗忘/输出门（范围0-1），tanh用于细胞门（范围-1-1）
        visual_i_gate = F.sigmoid(visual_i_gate)
        visual_f_gate = F.sigmoid(visual_f_gate)
        visual_c_gate = F.tanh(visual_c_gate)
        visual_o_gate = F.sigmoid(visual_o_gate)

        # 更新视觉细胞状态和隐藏状态（输出）
        # 细胞状态 = 遗忘门*上一时刻细胞状态 + 输入门*当前细胞门
        visual_cell = visual_f_gate * visual_cell + visual_i_gate * visual_c_gate
        # 隐藏状态 = 输出门*tanh(细胞状态)
        visual_output = visual_o_gate * torch.tanh(visual_cell)

        # 音频模态门控激活函数（同视觉模态）
        audio_i_gate = F.sigmoid(audio_i_gate)
        audio_f_gate = F.sigmoid(audio_f_gate)
        audio_c_gate = F.tanh(audio_c_gate)
        audio_o_gate = F.sigmoid(audio_o_gate)

        # 更新音频细胞状态和隐藏状态（输出）
        audio_cell = audio_f_gate * audio_cell + audio_i_gate * audio_c_gate
        audio_output = audio_o_gate * torch.tanh(audio_cell)

        return visual_output, visual_cell, audio_output, audio_cell


class Dual_lstm(nn.Module):
    """
    双LSTM模型类：封装Dual_lstm_cell，实现对序列数据的处理
    支持视觉和音频模态的并行LSTM编码，可选双向LSTM（当前注释了反向部分）
    """
    def __init__(self):
        super(Dual_lstm, self).__init__()

        # 模型参数配置
        self.video_input_dim = 512    # 视觉输入特征维度
        self.video_fc_dim = 512       # 视觉特征全连接层输出维度（当前未使用）
        self.d_model = 256            # LSTM隐藏状态维度（与Dual_lstm_cell的hidden_dim对应）

        # 视觉特征全连接层（当前forward中未使用，预留用于特征维度调整）
        self.v_fc = nn.Linear(self.video_input_dim, self.video_fc_dim)

        # 实例化双LSTM单元（视觉输入512维，音频输入128维，隐藏层256维）
        self.LSTM_cell = Dual_lstm_cell(visual_input_dim=512, audio_input_dim=128, hidden_dim=256)
        # 预留反向LSTM单元（当前注释，用于双向LSTM）
        # self.LSTM_cell_r = Dual_lstm_cell(visual_input_dim=512, audio_input_dim=128, hidden_dim=256)

        # 激活函数和dropout层（防止过拟合）
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)  # dropout概率0.2


    def forward(self, audio_feature, visual_feature):
        """
        前向传播：处理序列特征，输出视觉和音频的LSTM编码结果
        Args:
            audio_feature: 音频序列特征 (batch_size, seq_len, audio_input_dim=128)
            visual_feature: 视觉序列特征 (batch_size, seq_len, visual_input_dim=512)
        Returns:
            audio_output: 音频LSTM编码结果 (batch_size, seq_len, hidden_dim=256)
            visual_output: 视觉LSTM编码结果 (batch_size, seq_len, hidden_dim=256)
        """
        # 音频和视觉输入特征（直接使用原始输入，未做额外处理）
        audio_rnn_input = audio_feature
        visual_rnn_input = visual_feature

        # 初始化隐藏状态和细胞状态（根据设备是否支持CUDA选择存储位置）
        # 视觉模态初始状态（正向）
        if torch.cuda.is_available():
            visual_hidden = Variable(torch.zeros(visual_rnn_input.size(0), self.d_model).cuda())  # (batch_size, d_model)
            visual_hidden_r = Variable(torch.zeros(visual_rnn_input.size(0), self.d_model).cuda())  # 反向LSTM初始隐藏状态（预留）
        else:
            visual_hidden = Variable(torch.zeros(visual_rnn_input.size(0), self.d_model))
            visual_hidden_r = Variable(torch.zeros(visual_rnn_input.size(0), self.d_model))

        # 视觉模态细胞状态（正向）
        if torch.cuda.is_available():
            visual_cell = Variable(torch.zeros(visual_rnn_input.size(0), self.d_model).cuda())
            visual_cell_r = Variable(torch.zeros(visual_rnn_input.size(0), self.d_model).cuda())  # 反向LSTM初始细胞状态（预留）
        else:
            visual_cell = Variable(torch.zeros(visual_rnn_input.size(0), self.d_model))
            visual_cell_r = Variable(torch.zeros(visual_rnn_input.size(0), self.d_model))

        # 音频模态初始状态（正向）
        if torch.cuda.is_available():
            audio_hidden = Variable(torch.zeros(audio_rnn_input.size(0), self.d_model).cuda())
            audio_hidden_r = Variable(torch.zeros(audio_rnn_input.size(0), self.d_model).cuda())  # 反向LSTM初始隐藏状态（预留）
        else:
            audio_hidden = Variable(torch.zeros(audio_rnn_input.size(0), self.d_model))
            audio_hidden_r = Variable(torch.zeros(audio_rnn_input.size(0), self.d_model))

        # 音频模态细胞状态（正向）
        if torch.cuda.is_available():
            audio_cell = Variable(torch.zeros(audio_rnn_input.size(0), self.d_model).cuda())
            audio_cell_r = Variable(torch.zeros(audio_rnn_input.size(0), self.d_model).cuda())  # 反向LSTM初始细胞状态（预留）
        else:
            audio_cell = Variable(torch.zeros(audio_rnn_input.size(0), self.d_model))
            audio_cell_r = Variable(torch.zeros(audio_rnn_input.size(0), self.d_model))

        # 存储输出的列表
        visual_output = []    # 正向视觉LSTM输出
        audio_output = []     # 正向音频LSTM输出
        visual_output_r = []  # 反向视觉LSTM输出（预留）
        audio_output_r = []   # 反向音频LSTM输出（预留）

        # 获取序列长度（视觉和音频序列长度相同）
        length = visual_rnn_input.size(1)

        # 将初始状态转换为double类型，与输入特征数据类型匹配（避免类型不匹配错误）
        visual_hidden = visual_hidden.double()
        visual_cell = visual_cell.double()
        audio_hidden = audio_hidden.double()
        audio_cell = audio_cell.double()
        visual_hidden_r = visual_hidden_r.double()
        visual_cell_r = visual_cell_r.double()
        audio_hidden_r = audio_hidden_r.double()
        audio_cell_r = audio_cell_r.double()

        # 正向LSTM：按时间步遍历序列
        for i in range(length):
            # 取出当前时间步的视觉和音频特征
            # visual_rnn_input[:,i,:]：(batch_size, visual_input_dim)
            # audio_rnn_input[:,i,:]：(batch_size, audio_input_dim)
            visual_hidden, visual_cell, audio_hidden, audio_cell = self.LSTM_cell(
                visual_rnn_input[:,i,:], visual_hidden, visual_cell,
                audio_rnn_input[:,i,:], audio_hidden, audio_cell
            )
            # 保存当前时间步的输出
            visual_output.append(visual_hidden)
            audio_output.append(audio_hidden)

        # 将列表转换为张量（batch_size, seq_len, hidden_dim）
        visual_output = torch.stack(visual_output, dim=1)
        audio_output = torch.stack(audio_output, dim=1)

        # 反向LSTM（当前注释，用于双向编码）
        # for i in range(length):
        #     # 从序列末尾开始遍历（反向）
        #     visual_hidden_r, visual_cell_r, audio_hidden_r, audio_cell_r = self.LSTM_cell_r(
        #         visual_rnn_input[:, length-1-i, :], visual_hidden_r, visual_cell_r,
        #         audio_rnn_input[:, length-1-i, :], audio_hidden_r, audio_cell_r
        #     )
        #     visual_output_r.append(visual_hidden_r)
        #     audio_output_r.append(audio_hidden_r)
        # # 反向输出需要翻转回原始序列顺序，并与正向输出拼接
        # visual_output_r = torch.stack(visual_output_r, dim=1)
        # visual_output_r = torch.flip(visual_output_r, dims=[1])
        # audio_output_r = torch.stack(audio_output_r, dim=1)
        # audio_output_r = torch.flip(audio_output_r, dims=[1])
        # visual_output = torch.cat((visual_output, visual_output_r), dim=2)  # 拼接正向和反向特征（维度翻倍）
        # audio_output = torch.cat((audio_output, audio_output_r), dim=2)

        return audio_output, visual_output


# 测试代码（验证模型能否正常运行）
# model = Dual_lstm()
# visual_feature = torch.randn(32, 10, 512)  # 模拟32个样本，10个时间步，512维视觉特征
# audio_feature = torch.randn(32, 10, 128)   # 模拟32个样本，10个时间步，128维音频特征
# model(audio_feature, visual_feature)  # 输出应为(audio_output, visual_output)，形状均为(32,10,256)